"""
物理参数
"""
from .. import unit, Q_
from pint import Quantity
from numpy import ndarray
import numpy as np
from .constant import EARTH_ROTATION_ANGULAR_VELOCITY, EARTH_RADIUS

def coriolis_parameter(latitude: Quantity | ndarray | float) -> Quantity:
    """计算科里奥利参数，需要传入纬度，如果没有单位，则默认为度。
    如果是 Numpy 数组，返回的也是相同 Shape 的 Numpy 数组。
    计算公式为：
        f = 2 * Omega * sin(latitude)
    其中 Omega 为地球自转角速度。

    Calculate the Coriolis parameter, which is a function of latitude.
    If the input latitude does not have a unit, it is assumed to be in degrees.
    The formula is:
        f = 2 * Omega * sin(latitude)
    where Omega is the angular velocity of the Earth.

    Parameters
    ----------
    latitude : Quantity | ndarray | float
        纬度，如果没有单位则认为是度。
        Latitude in degrees or radians.

    Returns
    -------
    Quantity
        科里奥利参数，单位为 rad/s。

        The Coriolis parameter corresponding to the input latitude.

    """
    if isinstance(latitude, (int, float, ndarray)):
        latitude = Q_(latitude, 'degree')
    elif not isinstance(latitude, Quantity):
        raise TypeError("latitude must be a Quantity, ndarray or float")

    latitude = latitude.to('radian')
    f = 2 * EARTH_ROTATION_ANGULAR_VELOCITY * np.sin(latitude)
    return f


def rossby_parameter(latitude: Quantity | ndarray | float) -> Quantity:
    """计算罗斯贝参数，需要传入纬度，如果没有单位，则默认为度。
    如果是 Numpy 数组，返回的也是相同 Shape 的 Numpy 数组。
    计算公式为：
        beta = 2 * Omega * cos(latitude) / a
    其中 Omega 为地球自转角速度，a 为地球半径。

    Calculate the rossby parameter, which is a function of latitude.
    If the input latitude does not have a unit, it is assumed to be in degrees.
    The formula is:
        beta = 2 * Omega * cos(latitude) / a
    where Omega is the angular velocity of the Earth and a is the radius of the Earth.

    Parameters
    ----------
    latitude : Quantity | ndarray | float
        纬度，如果没有单位则认为是度。
        Latitude in degrees or radians.

    Returns
    -------
    Quantity
        罗斯贝参数，单位为 rad/s/m。

        The Rossby parameter corresponding to the input latitude.

    """
    if isinstance(latitude, (int, float, ndarray)):
        latitude = Q_(latitude, 'degree')
    elif not isinstance(latitude, Quantity):
        raise TypeError("latitude must be a Quantity, ndarray or float")

    latitude = latitude.to('radian')
    b = 2 * EARTH_ROTATION_ANGULAR_VELOCITY * np.cos(latitude) / EARTH_RADIUS
    return b